"""
这个脚本提供单线程/多线程读取hbase中网格数据的方法。
注意：需要在服务器上运行hbase的thrift服务器：
/usr/local/hbase/bin/start-hbase.sh
/usr/local/hbase/bin/hbase-daemon.sh start thrift
"""
import json
from concurrent.futures import ThreadPoolExecutor
from concurrent.futures import wait
import happybase
import time
import datetime
import numpy
import numpy as np

class GridReader:
    """
    单线程或者多线程读取一块区域的网格数据。
    如果使用多线程读取，那么会将输入区域切成多个小区域，每个小区域的最大范围是100x100。使用示例：
    (1)让脚本根据最大线程数(thread_count)自动选择多线程还是单线程工作模式。这是最省事的方法，推荐采用这种方式。
    reader = GridReader("192.168.1.119", 10000, thread_count)
    read_grid_data(118.1, 24.4, 305, 101)
    当thread_count > 1时用多线程，否则用单线程。
    (2)显式指定脚本在单线程模式下工作（格子区域小于100*100时建议使用），不需要设置thread_count最大线程数
    reader = GridReader("192.168.1.119", 10000)
    reader.read_grid_data_one_thread(118.1, 24.4, 80, 100)
    即使设置了thread_count最大线程数，在调用read_grid_data_one_thread时也会被忽略。
    (3)显式指定脚本在多线程模式下工作（格子区域大于100*100时建议使用），需要设置thread_count最大线程数量参数
    reader = GridReader("192.168.1.119", 10000, 8)
    reader.read_grid_data_multi_thread(118.1, 24.4, 305, 101)
    第一行代码指定hbase数据库的host和端口，设定同时运行多少个线程读取不同的块。
    第二行代码开始读取(118.1, 24.4)开始的x方向上的305个格子和y方向的101个格子围成的一个区域。
    """
    def __init__(self, host, port, thread_count = 8):
        if not host:
            raise Exception("hbase host not set")
        self.host = host
        self.port = 10000
        if port is not None:
            self.port = port
        self.thread_count = thread_count

    def read_grid_data_with_coord(self, coord_x, coord_y, x_cc, y_cc, index):
        """
        读取一小块区域
        coord_x：起始坐标x，需按10对齐。
        coord_y: 起始坐标y，需按10对齐。
        x_cc: 横向长度，需按10对齐
        y_cc: 纵向长度，需按10对齐
        index: 预留未用，在调试多线程时用于打印任务id。
        """
        assert coord_x % 10 == 0 and coord_y % 10 == 0
        assert x_cc % 10 == 0 and y_cc % 10 == 0
        hbase_conn = happybase.Connection(host=self.host, port=self.port)
        print(f"开始从hbase({self.host}:{self.port})读取地图……")
        try:
            table = hbase_conn.table('mapdata')
            start_time = time.perf_counter()
            # 我们要读取的范围是 左下坐标(coord_x, coord_y), 右上坐标(coord_x+x_cc, coord_y+y_cc)围起来的矩形区域。为了增加读写效率
            # 写入hbase时每一row-key对应了一个10x10网格的数据，因此在hbase中row-key的x分量和y分量都已按10对齐，假设现在已经读取了x=10
            # 的一列地图网格那么接下去读x=20的一列地图网格就可以。
            # TODO 我们一次性读取了一整列，如果这里未来成为了性能瓶颈就需要加个循环分步读取y的数据（比如从y_cc+coord_y开始每次加100）。
            # 写入的时候是这样做的(row-key的高32位存网格坐标的x值，低32位存网格坐标的y值)：
            # table.put('{:016x}'.format((x << 32) + y), {
            #           'route:accessible': json.dumps(list_accessible),
            #           'route:weight': json.dumps(list_weights),
            #           'geo:xml_text': json.dumps(list_xml_text),
            #           'geo:rect': json.dumps(list_polygon_wkt)
            #           })
            # 在下面的循环中table.scan： hbase的row-key根据地图网格坐标(x,y)和写入时作相同处理hex(x << 32 + y) 即读网格坐标x坐标范围
            # [coord_x, coord_x+x_cc) 网格坐标y坐标范围 [coord_y, coord_y+y_cc)。即每次循环读取网格坐标x的一列大网格（每个大网格又
            # 包含10x10个小网格)。
            # 例如第一次读[(0, 0), (0, 10), (0, 20) ... (0, y_cc-1)]处的大网格,第二次读取[(10, 0), (10, 10), (10, 20) ...
            # (10, y_cc-1)]处的大网格，依次类推。以下图为例，每个'x','o'代表一个10x10的地图网格。
            # y+40  *     o      o      o
            # y+30  *     o      o      o
            # y+20  *     o      o      o
            # y+10  *     o      o      o
            # y+0   x+0   x+10   x+20   x+30
            list_col = []
            for x in range(0, x_cc//10):
                # row_key: 每次scan的范围是地图网格坐标始于x, y范围在(coord_y~y_cc+coord_y)内的y_cc行数据。
                one_col = table.scan(row_start='{:016x}'.format(((coord_x + x * 10) << 32) + coord_y),
                                     row_stop='{:016x}'.format(((coord_x + x * 10) << 32) + y_cc + coord_y),
                                     columns=['route'])
                grid_ten_col_with_y_cc_high = []
                # 在hbase中每条数据是这样存的：
                # "row-key":{ "geo:rect":Polygons[100], "geo:xml_text": 格子osm[100], "route:accessible":byte[100],
                #             "route:weight": byte[100] }
                for e in one_col:  # 此处应有y_cc//10条数据，遍历它
                    # print("row-key:", e[0].decode())  # row-key 地理信息：e[1][b'geo:rect'],e[1][b'geo:xml_text']
                    # 由于每条hbase的数据存储了10x10的数据，所以实质上存储了(x~x+10)的10列网格数据。用矩阵存一下。
                    route_accessible = np.array(json.loads(e[1][b'route:accessible']), dtype=np.uint8)  #
                    route_weight = np.array(json.loads(e[1][b'route:weight']), dtype=np.uint8)
                    # 把accessible矩阵（某个网格是否可到达）和weight矩阵（通过某个网格的代价）合并，以便将来在寻路的时候使用。
                    # 下面两行的作用是合并矩阵，把类似[1,3,5,7,9]和[2,4,6,7,8]两个数组变成[[1,2],[3,4],[5,6],[7,8],[9,10]]
                    n1 = np.vstack([route_accessible, route_weight])
                    n2 = n1.transpose()
                    assert(n2.shape == (100, 2))
                    # 每个cell都是10列乘10行的数据。把100x2的矩阵转成10x10x2的矩阵
                    grid_ten_col_with_y_cc_high.append(n2.reshape((10, 10, 2)))

                # 遍历了y_cc//10条数据的大网格，给他横向堆叠一下，组成一个大小为(10,v_cc)的大矩阵。（是的！是横向堆叠，很反直觉）
                vertical_10x_ycc = np.hstack(grid_ten_col_with_y_cc_high)
                # 把这个(10乘以v_cc)大小的数组保存到一个列表中
                list_col.append(vertical_10x_ycc)
            # 把一个个(10,v_cc)大小的数组纵向堆叠起来，就是最终的结果了。（是的！是纵向堆叠，很反直觉）
            grid_final = np.vstack(list_col)
            # print(grid_final.shape)
            # print(grid_final)
            return grid_final
        except Exception as e:
            print(e)
        finally:
            hbase_conn.close()

    def read_grid_data(self, lon: float, lat: float, x_cc: int, y_cc: int):
        p_lon, p_lat = round(float(lon), 1), round(float(lat), 1)
        px_cc = x_cc if x_cc % 10 == 0 else (x_cc + 10) // 10 * 10
        py_cc = y_cc if y_cc % 10 == 0 else (y_cc + 10) // 10 * 10
        print("读取范围:", p_lon, p_lat, "开始的长宽为:", px_cc, py_cc, "的区域")
        if self.thread_count > 1:
            return self.read_grid_data_multi_thread(p_lon, p_lat, px_cc, py_cc)
        else:
            return self.read_grid_data_one_thread(p_lon, p_lat, px_cc, py_cc)

    def read_grid_data_one_thread(self, p_lon: float, p_lat: float, px_cc: int, py_cc: int):
        """
        单线程读取一块区域的网格，由x_cc和y_cc指定x轴和y轴上的网格数量
        """
        # 以经纬度-180.00,-90.00为原点，计算它的格子坐标。经纬度的1度对应xy坐标轴上的100个单位。
        # 防止float的精度问题，round一下
        coord_x : int = round((p_lon+180.0)*100)
        coord_y : int = round((p_lat+90.0)*100)
        return self.read_grid_data_with_coord(coord_x, coord_y, px_cc, py_cc, 0)

    def read_grid_data_multi_thread(self, lon: float, lat: float, x_cc: int, y_cc: int):
        """
        多线程读取一块区域的网格，每个小块区域的大小是100x100.
        """
        # (1)根据经纬度，获取在笛卡尔坐标系中的起点和终点的x,y坐标。
        # 当前约定是经纬度的"(-180.00,-90.00)"对应笛卡尔坐标系中的原点(0,0)，步长是0.01度，即x范围是0~36000,y范围是0~18000
        x_coord_begin, y_coord_begin = round((lon+180.0) * 100), round((lat+90.0) * 100)
        x_coord_end, y_coord_end = x_coord_begin + x_cc, y_coord_begin + y_cc

        #print("x range:", x_b, x_e, "y range:", y_b, y_e)
        # (2)我们是多线程执行，一个任务处理100x100，计算需要多少个任务可以执行完成。
        x_site = []  # 元素: (起点x坐标, 终点x坐标, x轴长度)
        x, y = x_coord_begin, y_coord_begin
        while x < x_coord_end:
            tmp = x + 100
            if tmp >= x_coord_end:
                x_site.append((x, x_coord_end, x_coord_end-x))
            else:
                x_site.append((x, tmp, 100))
            x = tmp
        #print(x_site)
        y_site = [] # 元素: (起点y坐标, 终点y坐标, y轴长度)
        while y < y_coord_end:
            tmpy = y + 100
            if tmpy >= y_coord_end:
                y_site.append((y, y_coord_end, y_coord_end - y))
            else:
                y_site.append((y, tmpy, 100))
            y = tmpy
        #print(y_site)
        # (3) 组织线程参数和线程运行结果的数据，放到result_of_threads中，每个线程分配一个。
        result_of_threads = {}  # (起点x坐标, 起点y坐标, x轴长度, y轴长度) : {"index": 序号, "future": 类似句柄, "grid": 网格数据}
        cnt = 0
        for x in x_site: #[(11807, 11907, 100), (11907, 12007, 100), (12007, 12107, 100), (12107, 12112, 5)]
            for y in y_site: #[(2437, 2537, 100), (2537, 2538, 1)]
                result_of_threads[(x[0], y[0], x[2], y[2])] = {"index": cnt}
                cnt += 1
        print("开始读取，分为{0}块多线程读取".format(len(result_of_threads)))
        # (4) 执行并等待所有线程执行完毕
        with ThreadPoolExecutor(max_workers=self.thread_count) as executor:
            futures = []
            for k, v in result_of_threads.items():
                f = executor.submit(self.read_grid_data_with_coord, k[0], k[1], k[2], k[3], v["index"])
                futures.append(f)
                v["future"] = f
            # 等待所有任务完成
            wait(futures)
        # (5) 合并所有线程的结果。最终结果是一个x_cc乘以x_cc的二维数组。
        # 先取所有结果
        for k, v in result_of_threads.items():
            v["grid"] = v["future"].result()
            if v["grid"] is None:
                print(round((k[0]-18000)/100, 2), round((k[1]-9000)/100, 2), k[2], k[3], "读取结果为空")
                raise Exception("data corrupt")

        #print(result_of_threads)
        # 先用笨办法合并(逐个点赋值比较慢)，虽然矩阵操作更快但是因为存在不同维度的矩阵不好合并。
        print(datetime.datetime.now(), "开始合并")
        merged = numpy.empty((x_cc, y_cc, 2), dtype=np.uint8)
        #merged = [[(0, 0) for _ in range(y_cc)] for _ in range(x_cc)]  # 创建横向x_cc乘以纵向y_cc个元素的二维数组。
        # result_of_threads：(起点x坐标, 起点y坐标, x轴长度, y轴长度) : {"index": 序号, "future": 类似句柄, "grid": 网格数据}
        for k, v in result_of_threads.items():  # k例子(11807, 2437, 100, 100)
            x_i_grid, y_i_grid = k[0] - x_coord_begin, k[1] - y_coord_begin  # 二维数组的下标，必须从0,0开始。
            vi = v["grid"]
            if vi is None:
                print(round((k[0]-18000)/100, 2), round((k[1]-9000)/100, 2), k[2], k[3], "读取结果为空")
                raise Exception("data corrupt")
            for x, col in enumerate(vi):
                for y, data in enumerate(col):
                    merged[x_i_grid+x][y_i_grid+y] = data
        print(datetime.datetime.now(), "完成合并")
        return merged
        # print(merged)
        # 优化：尝试用矩阵
        # 先纵向堆叠，再横向堆叠
        # grids = np.empty((x_cc, y_cc, 2), dtype=np.uint8)
        # for x in x_site:
        #     count_x = x[2]
        #     cols = np.empty((count_x, y_cc,2), dtype=np.uint8)
        #     for y in y_site:
        #         arr1 = np.array(roa[(x[0], y[0], x[2], y[2])]["grid"])
        #         cols = np.vstack((arr1, cols))
        #     grids = np.hstack((cols, grids))
        # print(grids)

# if __name__ == "__main__":
#     gg = GridReader("192.168.1.119", 10000, 1)
#     gg.read_grid_data(112.0, 19.3, 100, 100)
#     gg.read_grid_data(112.1,19.4, 100, 100)
